import cv2
import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms
import time
import os


class CatDogDataset(Dataset):
    def __init__(self, cat_paths, dog_paths, transform=None):
        self.images = []  
        self.labels = [] 
        self.transform = transform
        
  
        for path in cat_paths:
            self.images.append(path)
            self.labels.append(0)  

        

        for path in dog_paths:
            self.images.append(path)
            self.labels.append(1)  


    def __len__(self):
        return len(self.images)  

    def __getitem__(self, idx):
        
        img_path = self.images[idx]
        img = cv2.imread(img_path)

        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        img = transforms.ToPILImage()(img)
        if self.transform:
            img = self.transform(img)
        return img, self.labels[idx]  

class MiniCNN(nn.Module):
    def __init__(self):
        super(MiniCNN, self).__init__()
        self.features = nn.Sequential(
            nn.Conv2d(3, 32, 3, padding=1), 
            nn.ReLU(),
            nn.MaxPool2d(2, 2),  
            nn.Conv2d(32, 64, 3, padding=1), #
            nn.ReLU(), #
            nn.MaxPool2d(2, 2) #
        )
        
        self.classifier = nn.Sequential(
            nn.Flatten(),  
            nn.Linear(28*28*64, 256),  
            nn.ReLU(),
            nn.Linear(256, 2)
        )

    def forward(self, x):
        x = self.features(x)
        x = self.classifier(x)
        return x

train_transform = transforms.Compose([
    transforms.Resize((128, 128)),   
    transforms.RandomCrop(112),     
    transforms.RandomHorizontalFlip(p=0.5),  
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])  
])


test_transform = transforms.Compose([
    transforms.Resize((128, 128)),
    transforms.CenterCrop(112),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])


cat_paths = [f"./dataset/Cat/{i}.jpg" for i in range(0, 101)]
dog_paths = [f"./dataset/Dog/{i}.jpg" for i in range(0, 101)]


train_dataset = CatDogDataset(cat_paths, dog_paths, train_transform)
train_loader = DataLoader(train_dataset, batch_size=8, shuffle=True)


device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = MiniCNN().to(device)
criterion = nn.CrossEntropyLoss() 
optimizer = optim.Adam(model.parameters(), lr=0.001) 


epochs = 25
start_time = time.time()

print(f"在{device}上训练（100张/类样本，共{epochs}轮）...")
for epoch in range(epochs):
    model.train()
    train_loss, correct, total = 0.0, 0, 0
    for inputs, labels in train_loader:

        valid_mask = labels != -1
        if not valid_mask.any():
            continue 
        inputs = inputs[valid_mask].to(device)
        labels = labels[valid_mask].to(device)
        

        outputs = model(inputs)
        loss = criterion(outputs, labels)

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        

        train_loss += loss.item() * inputs.size(0)
        _, pred = torch.max(outputs, 1)
        total += labels.size(0)
        correct += (pred == labels).sum().item()
    

    epoch_loss = train_loss / total if total > 0 else 0.0
    epoch_acc = correct / total * 100 if total > 0 else 0.0
    print(f"第{epoch+1:2d}轮 | 损失：{epoch_loss:.4f} | 训练准确率：{epoch_acc:.2f}%")

print(f"\n训练完成！总耗时：{(time.time()-start_time)/60:.2f}分钟")


test_cat_dir = r"./dataset/test/Cat"
test_dog_dir = r"./dataset/test/Dog"


image_extensions = ('.jpg', '.jpeg', '.png', '.bmp', '.gif')


test_files = []


for filename in os.listdir(test_cat_dir):
    if filename.lower().endswith(image_extensions):
        file_path = os.path.join(test_cat_dir, filename)
        test_files.append((file_path, 0))  



for filename in os.listdir(test_dog_dir):
    if filename.lower().endswith(image_extensions):
        file_path = os.path.join(test_dog_dir, filename)
        test_files.append((file_path, 1))  



if not test_files:
    print("\n未找到任何测试图片！")
else:
    model.eval()
    correct = 0
    total = len(test_files)
    
    print(f"\n迷你CNN测试结果（共{total}张测试图）：")
    print("-" * 80)
    print(f"{'测试文件':<50} {'预测结果':<10} {'真实结果':<10} {'是否正确'}")
    print("-" * 80)
    
    with torch.no_grad():  
        for file_path, true_label in test_files:
            try:

                img = cv2.imread(file_path)
                
                img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
                img = transforms.ToPILImage()(img)
                img = test_transform(img).unsqueeze(0).to(device)

                outputs = model(img)
                _, pred = torch.max(outputs, 1)
                pred_label = pred.item()

                pred_res = "猫" if pred_label == 0 else "狗"
                true_res = "猫" if true_label == 0 else "狗"
                is_correct = "是" if pred_label == true_label else "否"
                
                if is_correct == "是":
                    correct += 1

                filename = os.path.basename(file_path)
                print(f"{filename:<50} {pred_res:<10} {true_res:<10} {is_correct}")
                
            except Exception as e:
                filename = os.path.basename(file_path)
                print(f"{filename:<50} 处理错误   未知       否 (错误: {str(e)[:20]})")
    

    test_acc = correct / total * 100
    print("-" * 80)
    print(f"测试总结：共测试{total}张图片，正确{correct}张，准确率：{test_acc:.2f}%")
